import numpy as np
from PIL import Image
import os
import torch


def write_rgb(tersor_img, epoch=None, file_list=None, file_name=None):
    b, c, h, w = tersor_img.shape

    file_name = file_name.replace('output_dir', 'main_result')
    path = './' + file_name + '/' + str(epoch) + '_epoch/' + str(epoch) + '/'
    if not os.path.exists(path):
        os.makedirs(path)
    imgs = torch.einsum('bchw->bhwc', tersor_img).detach().cpu().numpy()
    for i in range(b):
        img = imgs[i, :, :, :]
        img = Image.fromarray((img * 255).astype('uint8'))
        img.save(path + file_list)



def write_real_road_norm(tersor_img,pre_ir, pre_vis,epoch=None,file_list=None,file_name=None):
    b,c,h,w=tersor_img.shape

    file_name=file_name.replace('output_dir','main_result')
    path = './' +file_name+'/'+ str(epoch) + '_epoch/' + str(epoch) + '_Road/'
    path_ir = './' +file_name+'/'+ str(epoch) + '_epoch/' + str(epoch) + '_Road_ir/'
    path_vis = './' +file_name+'/'+ str(epoch) + '_epoch/' + str(epoch) + '_Road_vis/'
    if not os.path.exists(path):
        os.makedirs(path)
        os.makedirs(path_vis)
        os.makedirs(path_ir)
    imgs=torch.einsum('bchw->bhwc',tersor_img).detach().cpu().numpy()
    ir=torch.einsum('bchw->bhwc',pre_ir).detach().cpu().numpy()
    vis=torch.einsum('bchw->bhwc',pre_vis).detach().cpu().numpy()
    for i in range(b):
        img=imgs[i,:,:,:]
        img=Image.fromarray((img*255).astype('uint8'))
        img.save(path+file_list)
        img=ir[i,:,:,:]
        # print(img.shape)

        img=Image.fromarray((img*255).astype('uint8')[:,:,0],mode="L")
        img.save(path_ir+file_list)
        
        img=vis[i,:,:,:]
        img=Image.fromarray((img*255).astype('uint8')[:,:,0],mode="L")
        img.save(path_vis+file_list)


def write_real_tno_norm(tersor_img,epoch=None,file_list=None,file_name=None):
    b,c,h,w=tersor_img.shape


    file_name = file_name.replace('output_dir', 'main_result')
    path = './' + file_name + '/' + str(epoch) + '_epoch/' + str(epoch) + '_TNO/'
    if not os.path.exists(path):
        os.makedirs(path)
    imgs = torch.einsum('bchw->bhwc', tersor_img).detach().cpu().numpy()

    for i in range(b):
        img = imgs[i, :, :, :]
        img = Image.fromarray((img * 255).astype('uint8'))
        img.save(path + file_list)

def write_real_fake_norm(tersor_img,epoch=None,file_list=None,file_name=None):
    b,c,h,w=tersor_img.shape
    half_b=int(b/2)

    if file_name==None:
        path = './' + str(epoch) + '_epoch/' + str(epoch) + '_real/'
        if not os.path.exists(path):
            os.makedirs(path)
        path1 = './' + str(epoch) + '_epoch/' + str(epoch) + '_fake/'
        if not os.path.exists(path1):
            os.makedirs(path1)
    else:
        file_name=file_name.replace('output_dir','main_result')
        path = './' +file_name+'/'+ str(epoch) + '_epoch/' + str(epoch) + '_real/'
        if not os.path.exists(path):
            os.makedirs(path)
        path1 = './' +file_name+'/'+ str(epoch) + '_epoch/' + str(epoch) + '_fake/'
        if not os.path.exists(path1):
            os.makedirs(path1)

    imgs=torch.einsum('bchw->bhwc',tersor_img).detach().cpu().numpy()
    for i in range(half_b):
        img=imgs[i,:,:,:]
        img=Image.fromarray((img*255).astype('uint8'))
        img.save(path+file_list)
    for i in range(half_b,b):
        img=imgs[i,:,:,:]
        img=Image.fromarray((img*255).astype('uint8'))
        img.save(path1+file_list)
def write_real_fake(tersor_img,epoch=None,file_list=None):
    b,c,h,w=tersor_img.shape
    half_b=int(b/2)

    path='./'+str(epoch)+'_real/'
    if not os.path.exists(path):
        os.makedirs(path)
    path1='./'+str(epoch)+'_fake/'
    if not os.path.exists(path1):
        os.makedirs(path1)
    imgs=torch.einsum('bchw->bhwc',tersor_img).detach().cpu().numpy()
    for i in range(half_b):
        img=imgs[i,:,:,:]
        img=Image.fromarray((img*255).astype('uint8'))
        img.save(path+file_list)
    for i in range(half_b,b):
        img=imgs[i,:,:,:]
        img=Image.fromarray((img*255).astype('uint8'))
        img.save(path1+file_list)


def norm_to_225(tersor_img):
    device=tersor_img.device
    tersor_img=torch.einsum('bchw->bhwc',tersor_img)
    # imgs=tersor_img.detach().cpu().numpy()
    imgs=torch.tensor(np.array([0.229, 0.224, 0.225])).to(device) * tersor_img + torch.tensor(np.array([0.485, 0.456, 0.406])).to(device)
    imgs = torch.einsum('bhwc->bchw', imgs)
    return imgs
def write_tensor_img(tersor_img,file_name=None):
    tersor_img=torch.einsum('bchw->bhwc',tersor_img)
    path='./'+file_name+'/'
    if not os.path.exists(path):
        os.makedirs(path)
    imgs=tersor_img.detach().cpu().numpy()
    b,h,w,c=imgs.shape
    for i in range(b):
        img=imgs[i,:,:,:]
        img=Image.fromarray((img*255).astype('uint8'))
        img.save(path+str(i)+'.jpg')

def write_tensor_normal_img(tersor_img,file_name=None):
    tersor_img=torch.einsum('bchw->bhwc',tersor_img)
    path='./'+file_name+'/'
    if not os.path.exists(path):
        os.makedirs(path)
    imgs=tersor_img.detach().cpu().numpy()
    imgs=np.array([0.229, 0.224, 0.225]) * imgs + np.array([0.485, 0.456, 0.406])
    b,h,w,c=imgs.shape
    for i in range(b):
        img=imgs[i,:,:,:]
        img=Image.fromarray((img*255).astype('uint8'))
        img.save(path+str(i)+'.jpg')


def write_numpy_img(tersor_img,file_name=None):
    path='./'+file_name+'/'
    if not os.path.exists(path):
        os.makedirs(path)
    imgs=tersor_img
    b,h,w,c=imgs.shape
    for i in range(b):
        img=imgs[i,:,:,:]
        img=Image.fromarray((img*255).astype('uint8'))
        img.save(path+str(i)+'.jpg')